###################################################################
# To convert cat dataset .mat file to xml file
# and it is also to obtain images info,such as img weight and height
# author: greensroll
###################################################################


import mat4py
import os
import pandas as pd
from PIL import Image
from xml.dom.minidom import Document
from collections import OrderedDict

# parse the car dataset mat file
def parse_mat(filename):
    data_raw = mat4py.loadmat(filename)  # data_raw is a dict
    data = data_raw.get('annotations')  # data also a dict
    bbox_x1 = data.get('bbox_x1')
    bbox_y1 = data.get('bbox_y1')
    bbox_x2 = data.get('bbox_x2')
    bbox_y2 = data.get('bbox_y2')
    fname = data.get('fname')
    ann_info = OrderedDict([('fname', fname), ('xmin', bbox_x1), ('xmax', bbox_x2), ('ymin', bbox_y1), ('ymax', bbox_y2)])
    pd_info = pd.DataFrame(ann_info)
    return pd_info

def get_image_size(img_folder):
    img_info = dict()
    assert os.path.exists(img_folder)
    assert os.path.isdir(img_folder)
    imageList = os.listdir(img_folder)

    # read jpg file get's the width adn height
    for item in imageList:
        im = Image.open(os.path.join(img_folder, item))
        img_info[item] = [im.width, im.height]
    return img_info


def writeInfoToXml(xml_folder, pd_matdata, img_info):
    # 依次将matdata中的每一组元素提取出来，创建对应节点并插入dom树
    for i in range(pd_matdata.shape[0]):
        # 创建dom文档
        doc = Document()
        # 创建根节点
        annotation = doc.createElement('annotation')
        # 根节点插入dom树
        doc.appendChild(annotation)

        # 每一组信息先创建节点<filename>，然后插入到父节点<annotation>下
        filename = doc.createElement('filename')
        annotation.appendChild(filename)
        fname = doc.createTextNode(str(pd_matdata.iloc[i, :][0]))  # 元素内容写入
        filename.appendChild(fname)

        # 创建size节点，然后插入到父节点annotation下
        size = doc.createElement('size')
        annotation.appendChild(size)

        # 创建width节点，然后插入到父节点size下
        width = doc.createElement('width')
        img_width = doc.createTextNode(str(img_info[pd_matdata.iloc[i, :][0]][0]))
        width.appendChild(img_width)
        size.appendChild(width)

        # 创建height节点，然后插入到父节点size下
        height = doc.createElement('height')
        img_height = doc.createTextNode(str(img_info[pd_matdata.iloc[i, :][0]][1]))
        height.appendChild(img_height)
        size.appendChild(height)

        # 创建object节点，然后插入到父节点annotation下
        object_ = doc.createElement('object')
        annotation.appendChild(object_)

        # 创建name节点，然后插入到父节点object下
        name = doc.createElement('name')
        obj_name = doc.createTextNode('car')
        name.appendChild(obj_name)
        object_.appendChild(name)

        # 创建bndbox节点，然后插入到父节点object下
        bndbox = doc.createElement('bndbox')
        object_.appendChild(bndbox)

        # 创建xmin, xmax, ymin, ymax节点，然后插入到父节点badbox下
        xmin_ = doc.createElement('xmin')
        value_xmin = doc.createTextNode(str(pd_matdata.iloc[i, :][1]))
        xmin_.appendChild(value_xmin)
        bndbox.appendChild(xmin_)

        xmax_ = doc.createElement('xmax')
        value_xmax = doc.createTextNode(str(pd_matdata.iloc[i, :][2]))
        xmax_.appendChild(value_xmax)
        bndbox.appendChild(xmax_)

        ymin_ = doc.createElement('ymin')
        value_ymin = doc.createTextNode(str(pd_matdata.iloc[i, :][3]))
        ymin_.appendChild(value_ymin)
        bndbox.appendChild(ymin_)

        ymax_ = doc.createElement('ymax')
        value_ymax = doc.createTextNode(str(pd_matdata.iloc[i, :][4]))
        ymax_.appendChild(value_ymax)
        bndbox.appendChild(ymax_)

        # 将dom对象写入本地xml文件
        # print(fname.split('.')[0])
        xml_filename = os.path.join(xml_folder, str(pd_matdata.iloc[i, :][0]).split('.')[0] + '.xml')
        #txt_filename = os.path.join('/home/greensroll/DeepLearning/project/car_dataset', 'trainval' + '.txt')
        #with open(txt_filename, 'a') as f:
        #    f.write(str(pd_matdata.iloc[i, :][0]).split('.')[0])
        #    f.write('\n')
        with open(xml_filename, 'wt') as f:
            f.write(doc.toprettyxml(indent='\t', encoding='utf-8'))
            doc.writexml(f, indent='\t', newl='\n', addindent='\t', encoding='utf-8')
    return

if __name__ == '__main__':
    filename = '/devkit/cars_train_annos.mat'
    xml_folder = '/annotation'
    img_folder = '/cars_train'

    pd_info = parse_mat(filename)
    img_info = get_image_size(img_folder)
    writeInfoToXml(xml_folder, pd_info, img_info)